# -- coding: utf-8 --
"""
压力修正方程求解
"""

# third party lib
import numpy as np

# private lib
from utils.constant import InterpMode
from utils.numeric import grid_interp, grid_extrap, difference, thomas_algorithm



def cal_pressure_prime(params: dict, liquid_vapour_mass_transfer_rate:np.ndarray) -> np.ndarray:
    """ 压力修正方程计算，用三对角矩阵求解
    Args:
        params: 计算需要的输入参数
        liquid_vapour_mass_transfer_rate: np.ndarray 液相蒸气质量传递速率

    Returns:
        p_prime_node_LHS: np.ndarray, 压力误差修正量
    """
    dx = params["dx"]
    dt = params["dt"]
    rho_g_node = params["rho_g_node"]
    rho_L_node = params["rho_L_node"]
    rho_g_segment = grid_interp(rho_g_node, InterpMode.MID)
    rho_L_segment = grid_interp(rho_L_node, InterpMode.MID)
    alpha_g_node = params["alpha_g_node"]
    alpha_L_node = params["alpha_L_node"]
    alpha_D_node = params["alpha_D_node"]
    alpha_g_segment = grid_interp(alpha_g_node, InterpMode.MID)
    alpha_L_segment = grid_interp(alpha_L_node, InterpMode.MID)
    alpha_D_segment = grid_interp(alpha_D_node, InterpMode.MID)
    
    c_g_node = params["c_g_node"]
    c_L_node = params["c_L_node"]
    
    rho_g_node_iteration_init = params["rho_g_node_iteration_init"]
    rho_L_node_iteration_init = params["rho_L_node_iteration_init"]
    
    V_g_segment = params["V_g_segment"]
    V_L_segment = params["V_L_segment"]
    V_D_segment = params["V_D_segment"]
    
    # 构造三对角矩阵系数
    a_E = - 1 / rho_g_node[1:-1] * dt / dx * alpha_g_segment[:-1] - 1 / rho_L_node[1:-1] * dt / dx * alpha_L_segment[:-1] - 1 / rho_L_node[1:-1] * dt / dx * alpha_D_segment[:-1]

    a_W = -1 / rho_g_node[1:-1] * dt / dx * alpha_g_segment[1:] - 1 / rho_L_node[1:-1] * dt / dx * alpha_L_segment[1:] - 1 / rho_L_node[1:-1] * dt / dx * alpha_D_segment[1:]

    # 主对角元素
    a_P = - a_W - a_E + dx / dt * (
            1 / c_g_node ** 2 * alpha_g_node / rho_g_node
            + 1 / c_L_node ** 2 * alpha_L_node / rho_L_node
            + 1 / c_L_node ** 2 * alpha_D_node / rho_L_node
        )[1:-1]
    
    
    psi_g_node = grid_extrap(liquid_vapour_mass_transfer_rate, 1)
    RHS_b_src = (psi_g_node / rho_g_node) + (-psi_g_node / rho_L_node)

    # 右侧向量
    RHS_b = (
        -dx / dt * (
            (alpha_g_node / rho_g_node) 
            * (rho_g_node - rho_g_node_iteration_init)
            + (alpha_L_node / rho_L_node)
            * (rho_L_node - rho_L_node_iteration_init)
            + (alpha_D_node / rho_L_node)
            * (rho_L_node - rho_L_node_iteration_init)
        )[1:-1]
        - 1 / rho_g_node[1:-1]
        * difference(alpha_g_segment * rho_g_segment * V_g_segment, dx, 1)
        * dx
        - 1 / rho_L_node[1:-1]
        * difference(alpha_L_segment * rho_L_segment * V_L_segment, dx, 1)
        * dx
        - 1 / rho_L_node[1:-1]
        * difference(alpha_D_segment * rho_L_segment * V_D_segment, dx, 1)
        * dx
        + RHS_b_src[1:-1]
    )
    
    # 边界修正
    RHS_b[0] -= a_W[0] * RHS_b_src[0]
    RHS_b[-1] -= a_E[-1] * RHS_b_src[-1]

    # 求解三对角矩阵方程
    p_prime_inner = thomas_algorithm(
        a_W=a_W,  # 下对角
        a_P=a_P,  # 主对角
        a_E=a_E,  # 上对角
        b=RHS_b # 右侧向量
    )
    
    # 构造完整的压力修正量数组(包含边界)
    p_prime_node = np.zeros_like(params["p_node"])
    p_prime_node[1:-1] = p_prime_inner
    return p_prime_node
